import torch
import torch.nn as nn
import torch.optim as optim

# 定义一个简单模型
model = nn.Linear(10, 2)

# 定义损失函数和优化器
loss_fn  = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)

# 模拟输入和标签
input = torch.randn(4, 10) # batch size 4, 10 features
print("Input shape:", input.shape)
print(input)
labels = torch.tensor([0, 1, 0, 1]) # batch size 4, binary labels

for epoch in range(100):
    # 前向传播
    outputs = model(input)
   
    # 计算损失
    loss = loss_fn(outputs, labels)

    # 反向传播和优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

# 获取模型的状态字典
print("Model state dict:")
print(model.state_dict())
